import dv_processing as dv
import cv2
import numpy as np
import os
import glob
import argparse
from tqdm import tqdm
from metavision_core.event_io import EventsIterator

np.random.seed(0)

def parse_argument():
    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument('--raw_file_dir', type=str, default='E:/Datasets/uav_event&rgb', help='The path of a raw dir')
    parser.add_argument('--save_dir', type=str, default='E:/Datasets/uav_eventframe', help='')
    parser.add_argument('--time_interval_us', type=int, default=50000, help='Time interval in microseconds')
    parser.add_argument('--show', default=True, help='If show detection and tracking process')
    return parser

def get_reader(file_path, time_interval_us):
    assert os.path.exists(file_path), f"The file '{file_path}' does not exist"
    mv_iterator = EventsIterator(input_path=file_path, delta_t=time_interval_us)
    return mv_iterator

def process_events(reader, save_path, time_interval_us):
    all_events = []
    total_time = 0
    for evs in reader:
        all_events.append(evs)
        if len(evs) > 0:
            total_time = max(total_time, np.max(evs['t']))
    
    end_time = total_time - 50000  # 去掉末尾50ms
    start_time = max(0, end_time - 10000000)  # 只取最后10s

    # 按时间间隔处理事件
    height, width = reader.get_size()
    ev_cnt = 0
    current_time = start_time
    
    while current_time < end_time:
        event_image = np.zeros((height, width))
        next_time = min(current_time + time_interval_us, end_time)
        
        # 合并指定时间范围内的事件
        for evs in all_events:
            mask = (evs['t'] >= current_time) & (evs['t'] < next_time)
            if np.any(mask):
                np.add.at(event_image, (evs['y'][mask], evs['x'][mask]), 1)
        
        if np.sum(event_image) > 0:  # 只保存有事件的帧
            cv2.imwrite(f'{save_path}/{ev_cnt:06d}_{current_time}_{next_time}.png', 
                       (event_image * 60).astype(np.uint8))
            ev_cnt += 1
        
        current_time = next_time

if __name__ == '__main__':
    args, _ = parse_argument().parse_known_args(None)
    print(args)

    raw_path_list = []
    for root, dirs, files in os.walk(args.raw_file_dir):
        for file in files:
            if file.endswith('.raw'):
                raw_path_list.append(os.path.join(root, file))

    print(raw_path_list)

    for raw_path in tqdm(raw_path_list):
        assert os.path.exists(raw_path)
        base_name = os.path.basename(raw_path).split('.')[0]
        print(base_name)

        ev_save_path = os.path.join(args.save_dir, base_name, 'event')
        if not os.path.exists(ev_save_path):
            os.makedirs(ev_save_path)

        reader = get_reader(raw_path, args.time_interval_us)
        process_events(reader, ev_save_path, args.time_interval_us)